import os
import sys
import subprocess  
import shutil  

from transformers import AutoModel, AutoTokenizer

from langchain_community.chat_models import ChatZhipuAI
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langchain.output_parsers import StructuredOutputParser, ResponseSchema
from langchain_core.prompts import PromptTemplate

from docx import Document
from docx.shared import Pt
from docx.enum.text import WD_ALIGN_PARAGRAPH
from docx.image.exceptions import UnrecognizedImageError  
from PIL import Image
import json

tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, device_map='cuda', use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
model = model.eval().cuda()

#降低视频帧率为17,生成cut.mp4
def convert_video_framerate(input_video):  
    output_video='cut.mp4'
    framerate=17

    # 构建ffmpeg命令  
    command = [  
        'ffmpeg',  
        '-i', input_video,  # 输入文件  
        '-r', str(framerate),    # 设置帧率  
        output_video        # 输出文件  
    ]  
    # 执行命令  
    try:  
        # 使用subprocess.run执行命令，捕获输出和错误  
        print("正在使用ffmpeg转换视频文件的帧率(default=17)。")
        result = subprocess.run(command, capture_output=True, text=True, check=True)  
    except subprocess.CalledProcessError as e:  
        # 如果ffmpeg命令执行失败，则打印错误信息  
        print(f"ffmpeg命令执行失败，错误：")  
        print(e.stderr)  
    extract_keyframes(output_video)

#对cut.mp4关键帧提取，生成关键帧目录cut_ppt
def extract_keyframes(video_path):  
    # 构建输出文件夹路径  
    output_dir = os.path.join(os.path.dirname(video_path), "cut_ppt")   
    if not os.path.exists(output_dir):  
        os.makedirs(output_dir)  
    # 构建输出文件的路径模板  
    output_pattern = os.path.join(output_dir, "cut_%03d.jpg")  
      
    # 构建ffmpeg命令  
    command = [  
        'ffmpeg',  
        '-i', video_path,  
        '-vf', 'select=eq(pict_type\,I)',  
        '-vsync', 'vfr',  
        output_pattern  
    ]  
      
    # 执行ffmpeg命令  
    try:  
        print("正在进行关键帧提取...")
        subprocess.run(command, check=True)  
        print(f"关键帧已成功提取到 {output_dir}")  
    except subprocess.CalledProcessError as e:   
        print(f"命令执行失败，错误：")  
        print(e.stderr) 
    extract_keyframe_times()

#提取视频关键帧的时间戳,保存到time.txt
def extract_keyframe_times(video_file='cut.mp4', output_file='time.txt'):
    """
    使用 ffprobe 提取视频关键帧的时间戳，并将结果保存到指定的输出文件中。
    :param video_file: 输入视频文件路径
    :param output_file: 输出时间戳文件路径
    """
    # 构建 ffprobe 命令
    command = [
        'ffprobe',
        '-v', 'error',
        '-skip_frame', 'nokey',
        '-select_streams', 'v:0',
        '-show_entries', 'frame=pts_time',
        '-of', 'csv=print_section=0',
        video_file
    ]
    
    try:
        # 执行命令并将输出重定向到文件
        with open(output_file, 'w') as f:
            result = subprocess.run(command, stdout=f, check=True)
    
        print(f"关键帧时间戳已成功提取并保存到 {output_file}")
    except subprocess.CalledProcessError as e:
        print(f"执行 ffprobe 命令时发生错误: {e}")
    except Exception as e:
        print(f"发生未知错误: {e}")
    format_and_rename_images()

#重命名cut_ppt下的图片文件，方便生成json中的时间戳
def format_and_rename_images(image_dir='cut_ppt', time_file='time.txt'):
    """
    读取时间文件，格式化时间并重命名指定目录下的图片文件。
    
    :param image_dir: 图片文件所在的目录
    :param time_file: 包含时间戳的文本文件
    """
    # 读取原始时间数据
    with open(time_file, 'r') as f:
        times = f.readlines()

    # 处理每一行的数据
    formatted_times = []  
    for time in times:  
        # 去除字符串两端的空白字符  
        trimmed_time = time.strip()  
        # 检查字符串是否为空或只包含空格  
        if trimmed_time:  
            # 尝试将字符串转换为浮点数  
            try:  
                seconds = int(float(trimmed_time))  
            except ValueError:  
                # 如果转换失败，则打印错误信息并跳过当前项  
                print(f"Error converting '{trimmed_time}' to float, skipping.")  
                continue  
        else:  
            # 如果字符串为空或只包含空格，则打印警告并跳过当前项  
            print(f"Skipping empty or whitespace-only time '{time}'")  
            continue  
        # 计算分钟和秒  
        minutes = seconds // 60  
        remaining_seconds = seconds % 60  
        # 格式化为分:秒  
        formatted_time = f"{minutes}_{remaining_seconds:02d}"  
        formatted_times.append(formatted_time)  
  
    # 现在 formatted_times 包含了所有成功转换并格式化的时间
    # 将格式化后的时间写回到原文件
    with open(time_file, 'w') as f:
        f.write("\n".join(formatted_times) + "\n")

    # 读取时间文件中的内容
    with open(time_file, 'r') as f:
        times = f.read().splitlines()

    # 获取目录下的所有图片文件
    image_files = sorted([f for f in os.listdir(image_dir) if f.endswith('.jpg')])

    # 重命名图片文件
    for i, image_file in enumerate(image_files):
        if i < len(times):
            # 生成新的文件名
            new_name = f"{str(i + 1).zfill(3)}_{times[i]}.jpg"
            # 构建完整的旧路径和新路径
            old_path = os.path.join(image_dir, image_file)
            new_path = os.path.join(image_dir, new_name)
            # 重命名文件
            os.rename(old_path, new_path)




#识别所有关键帧,输出output.txt
def process_images_in_directory(image_dir='cut_ppt', output_txt='output.txt'):  
    """  
    遍历指定目录中的所有图片文件（.png 和 .jpg），使用GOT-OCR识别其内容，并将结果写入到指定的输出文件中。  
  
    参数:  
    image_dir (str): 包含图片文件的目录路径。  
    output_txt (str): 输出文件的名称，默认为'output.txt'。  
    """  
    print("正在OCR关键帧内容")
    # 假设 model 和 tokenizer 已经在外部定义  
    with open(output_txt, 'w', encoding='utf-8') as f:  
        cnt = 1
        # 遍历目录中的所有文件  
        for filename in sorted(os.listdir(image_dir)):  
            # 检查文件扩展名  
           if filename.lower().endswith(('.png', '.jpg')):  
                # 构建完整的文件路径  
                image_file = os.path.join(image_dir, filename)  
                # 执行 OCR 识别   
                res = model.chat(tokenizer, image_file, ocr_type='ocr')  
                
                parts = filename.split('_')
                if len(parts) > 2 :
                    timestamp = parts[1] + '_' + parts[2].split('.')[0] 
                else:
                    timestamp = "NaN"

                f.write(f"{timestamp}: {res}\n\n") 
                print(f"识别并写入关键帧{cnt},时间：{timestamp}") 
                cnt += 1
    print("所有关键帧内容已存放在output.txt")


#智谱轻言提取习题内容,  output.txt -> output.json
def extract_and_save_questions(api_key, txt_file_path, output_json_path):
    print("正在使用大模型提取结构化习题内容")
    # 设置API密钥
    os.environ["ZHIPUAI_API_KEY"] = api_key

    # 初始化模型
    llm = ChatZhipuAI(
        model="glm-4-plus",
        temperature=0.5,
    )

    # 解析格式
    response_schemas = [
        ResponseSchema(name="questions", description="A list of extracted questions, each as a structured object", properties=
        [
            ResponseSchema(name="text", description="The text of the question"),
            ResponseSchema(name="type", description="The type of question "),
            ResponseSchema(name="answer", description="The correct answer to the question"),
            ResponseSchema(name="timestamp", description="The timestamp of the question")
        ])
    ]

    output_parser = StructuredOutputParser.from_response_schemas(response_schemas)
    format_instructions = output_parser.get_format_instructions()

    # prompt
    summarizing_prompt_template = """
    {format_instructions}

    此文本是对课堂视频中文字内容的识别，每次识别以图片帧出现在视频中的时间戳作为开头（如"0分54秒:",但在输出时请不要带冒号）。
    接下来请提取文本中涉及到的地理、数学、物理、语文等习题。
    对于识别到的每道习题，输出格式为json的结构化文件，其中必须包括四个结构：问题内容(text)、类型(type)、答案(answer)和时间戳(timestamp)。
    如果内容(text)、类型(type)、答案(answer)和时间戳(timestamp)无法识别或者不确定，那么用空值代替即可，但不可以没有其中任何一个。
    ---
    {content}
    """
    prompt = PromptTemplate.from_template(summarizing_prompt_template, partial_variables={'format_instructions': format_instructions})

    summarizing_chain = prompt | llm | output_parser
    # 读取文本文件
    def read_txt(file_path):
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
        return content
    doc = read_txt(txt_file_path)

    # 提取问题并保存到JSON文件
    res = summarizing_chain.invoke({"content": doc})
    
    with open(output_json_path, 'w', encoding='utf-8') as f:
        json.dump(res, f, ensure_ascii=False)


# json 2 word
def create_word_document_from_json(json_file_path, images_dir, output_docx_path):
    print("正在将结构化习题转化为word")
    # 创建一个新的 Word 文档
    doc = Document()
    # 打开一个 JSON 文件
    with open(json_file_path, 'r', encoding='utf-8') as f:
        json_data = json.load(f)

    # 添加文档标题
    doc.add_heading('任务二：例题联系析取答案', 0)

    # 遍历 JSON 数据中的每个问题
    for question in json_data['questions']:
        # 问题
        if question['text']:
            doc.add_heading(f"问题: {question['text']}", level=1)
        # 类型
        if question['type']:
            doc.add_paragraph(f"类型: {question['type']}")
        # 答案
        if question['answer']:
            doc.add_paragraph(f"答案: {question['answer']}")

        if question['timestamp']:  
            doc.add_paragraph(f"时间戳: {question['timestamp']}")  
            # 匹配相应图片  
            timestamps = question['timestamp'].rstrip(':')  
            current_dir = os.path.abspath('.')
            for i in range(1, 1000):  
                prefix = f"{i:03d}"  
                image_filename = f"{prefix}_{timestamps}.jpg"  
                image_path = os.path.join(current_dir, images_dir, image_filename)  

                # 检查图片文件是否存在  
                if os.path.exists(image_path):  
                    try:  
                        # 在文档中插入图片  
                        doc.add_picture(image_path, width=Pt(400))  
                        break  # 找到并插入图片后退出循环  
                    except UnrecognizedImageError:  
                        # 如果图片无法识别，记录错误并继续尝试下一个文件  
                        print(f"无法识别图片文件 {image_path}。跳过...")  
                        continue 
                elif i == 999:  
                    paragraph = doc.add_paragraph("对应的图片未找到。")  
                    paragraph.alignment = WD_ALIGN_PARAGRAPH.CENTER  
                else:  
                    continue  # 如果当前图片不存在，则继续循环  
        else:  
            continue  

    # 保存文档
    doc.save(output_docx_path)
    return output_docx_path


if __name__ == "__main__":  
    if len(sys.argv) < 2:  
        print("请提供一个视频文件作为参数,例如：python xx.py xx.mp4")  
    else:  
        video_file = sys.argv[1]  # 获取第一个命令行参数，即视频文件名  

    convert_video_framerate(video_file)  
    process_images_in_directory()
 

    api_key = "34e4250edb1cc63dfde749a7c83fbbc1.2RbubnzAdoUcBmxM"
    extract_and_save_questions(api_key, 'output.txt', 'output.json') #从识别文本中提取结构化习题

    create_word_document_from_json('output.json', 'cut_ppt', 'output.docx')    #结构化习题转word，同时插入图片，增加可读性



    # 删除多余文件
    files_to_delete = ['time.txt', 'cut.mp4']
    for file_name in files_to_delete:
        try:
            os.remove(file_name)
        except OSError as e:
            print(f"部分文件清理失败，请手动删除： {file_name}: {e}")


    # 移动部分文件
    video_file_name, _ = os.path.splitext(os.path.basename(video_file))
    new_dir = os.path.join(os.getcwd(), video_file_name)
    if not os.path.exists(new_dir):
        os.makedirs(new_dir)

    files_to_move = ['output.txt', 'output.json', 'output.docx']
    for file_name in files_to_move:
        src = os.path.join(os.getcwd(), file_name)
        dst = os.path.join(new_dir, file_name)
        if os.path.isfile(src):
            shutil.move(src, dst)
        else:
            print(f"警告：{file_name} 不存在，无法移动。")

    # 移动cut_ppt目录中的所有jpg图片
    cut_ppt_dir = os.path.join(os.getcwd(), 'cut_ppt')
    if os.path.isdir(cut_ppt_dir):
        for file_name in os.listdir(cut_ppt_dir):
            if file_name.endswith('.jpg'):
                full_file_name = os.path.join(cut_ppt_dir, file_name)
                shutil.move(full_file_name, new_dir)
    else:
        print("cut_ppt目录不存在，无法移动图片。")

    print(f"已提取 {video_file} 中的习题，请前往文件夹查看！")

